from sklearn.feature_extraction.text import TfidfVectorizer
from typing import List

class KeywordExtractor:
    def __init__(self, domain_keywords: List[str] = None):
        self.domain_keywords = domain_keywords or []
        self.vectorizer = TfidfVectorizer(max_features=364)

    def extract(self, text: str, context: List[str] = None, top_n: int = 5) -> List[str]:
        """提取关键词（结合TF-IDF和领域词典）"""
        # 领域关键词优先
        domain_kws = [kw for kw in self.domain_keywords if kw in text]
        
        # TF-IDF提取通用关键词
        corpus = context + [text] if context else [text]
        tfidf_kws = self.vectorizer.fit_transform(corpus)
        feature_names = self.vectorizer.get_feature_names_out()
        scores = tfidf_kws[-1].toarray().flatten()
        top_indices = scores.argsort()[-top_n:][::-1]
        generic_kws = [feature_names[i] for i in top_indices]

        return list(set(domain_kws + generic_kws))
